import jieba
from collections import Counter
import re
import pdfplumber
import os
import time

from const import keywords, pdf_name

# 如果本地不存在 pdf 目录，则创建一个
if not os.path.exists(pdf_name):
    os.mkdir(pdf_name)

# 获取当前工作目录
current_directory = os.path.join(os.getcwd(), pdf_name)

# 列出当前目录下所有文件和子目录（不递归）
annual_reports_urls = os.listdir(current_directory)

# 存储每家公司关键词词频的字典
company_keyword_counts = {}

# 按年存储每家公司关键词词频的字典
year_keyword_counts = {}

# 存储所有年报中关键词总词频的计数器
total_keyword_counts = Counter()

start_time = time.time()
# 遍历每家公司的年报URL
for filename in annual_reports_urls:

    # 判断是否为pdf文件
    if not filename.endswith('.pdf'):
        print('%s==== 不是pdf文件 跳过' % filename)
        continue

    # print('%s====' % filename)
    # 公司名字
    company_name = re.search(r'(\D+)', filename).group(1)
    # 年份
    print('%s====' % filename)
    year = re.search(r'(\d+)', filename).group(1)
    # 文件绝对地址
    file_path = os.path.join(current_directory, filename)
    # 使用pdfplumber提取PDF中的文本
    text = ''

    print('%s==== 开始读取' % filename)
    with pdfplumber.open(file_path) as pdf:
        for page in pdf.pages:
            text += page.extract_text() or ''
        print('%s==== 结束读取' % filename)

    # 数据预处理，例如转换为小写
    # text = text.lower()
    # 英语单词：关键词词频统计
    # words = re.findall(r'\w+', text)  # 使用正则表达式匹配单词
    # 中文单词：使用jieba进行分词
    words = jieba.lcut(text)  # 对文本进行精确模式分词

    print('%s==== 正在计算关键词' % filename)
    word_counts = Counter(words)
    keyword_counts = { keyword: word_counts[keyword] for keyword in keywords }
    print('%s==== 结束计算关键词' % filename)

    # 将当前公司的关键词词频按年存储起来
    year_keyword_counts[company_name+year+"年"] = Counter(keyword_counts)

    # 将当前公司的关键词词频存储起来
    print('%s==== 正在关键词累计更新' % filename)
    if company_name not in company_keyword_counts:
        company_keyword_counts[company_name] = Counter(keyword_counts)
    else:
        company_keyword_counts[company_name].update(Counter(keyword_counts))
    print('%s==== 结束关键词累计更新' % filename)
    
    # 更新总词频计数器
    print('%s==== 正在为所有公司关键词累计更新' % filename)
    total_keyword_counts.update(Counter(keyword_counts))
    print('%s==== 结束为所有公司关键词累计更新' % filename)

end_time = time.time()
print('程序运行时间：%d秒' % (end_time - start_time))


# 输出公司每一年的关键词词频
print("每家公司的按年的关键词词频:")
for company_and_year, counts in year_keyword_counts.items():
    print(f"{company_and_year}: {counts}")
print("----------------------------------------")

# 输出每家公司的关键词词频
print("每家公司的关键词词频:")
for company, counts in company_keyword_counts.items():
    print(f"{company}: {counts}")
print("----------------------------------------")

# 输出所有年报中关键词的总词频
print("\n总的关键词词频:")
print(total_keyword_counts)
